import torch
from torch import nn

"""
TransformerDecoderLayer：
===========================__init___参数：=================================
d_model (int) –输入中的预期特征数 （必填）.
nhead (int) – MultiHeadAttention 模型中的 head 数（必需）.
dim_feedforward (int) – 前馈网络模型的维度(default=2048).
dropout (float) – dropout 值 (default=0.1).
activation (Union[str, Callable[[Tensor], Tensor]]) – 中间层的激活函数，可以是字符串 （“relu” 或 “gelu”） 或一元可调用对象. Default: relu
layer_norm_eps (float) – 图层归一化组件中的 EPS 值(default=1e-5).
batch_first (bool) – 如果 ，则提供输入和输出张量 as （batch， seq， feature） 的 AS （Batch， seq， feature） 中。True Default: False (seq, batch, feature).
norm_first (bool) –如果 ，层归一化在自我注意之前完成，则多头 attention 和 Feedforward作。否则，它将在True. Default: False (after).
bias (bool) – 如果设置为 ，则图层不会学习 Additive 偏置。FalseLinearLayerNormDefault: True.

TransformerDecoder
===========================__init___参数：=================================
decoder_layer (TransformerDecoderLayer) – TransformerDecoderLayer（） 类的实例（必需）.
num_layers (int) – 解码器中的子解码器层数（必填）.
norm (Optional[Module]) –层归一化组件（可选）.
===========================forward参数：=================================
tgt (Tensor) – 解码器的输入序列 （必需）。最后一个维度 embedding_dims
memory (Tensor) – 编码器的输出序列 （必需）。最后一个维度 embedding_dims
tgt_mask (Optional[Tensor]) – 单独序列掩码（填充掩码）
memory_mask (Optional[Tensor]) – 编码器的输出序列掩码
tgt_key_padding_mask (Optional[Tensor]) – 解码器的批次序列掩码（因果掩码）
memory_key_padding_mask (Optional[Tensor]) – 编码器的批次输出序列掩码（因果掩码）
tgt_is_causal (Optional[bool]) - 解码器是否是因果掩码
memory_is_causal (bool) - 编码器是否是因果掩码 （很少见）
"""

decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
memory = torch.rand(10, 32, 512)  # transformer_decoder的结果
tgt = torch.rand(20, 32, 512)
out = transformer_decoder(tgt, memory)
print(out.shape)
